{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 时间序列预测案例一: 正弦波\n",
    "\n",
    "> PyTorch 官方给出了时间序列的预测案例:\n",
    "https://github.com/pytorch/examples/tree/master/time_sequence_prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sequence(\n",
      "  (lstm1): LSTMCell(1, 51)\n",
      "  (lstm2): LSTMCell(51, 51)\n",
      "  (linear): Linear(in_features=51, out_features=1, bias=True)\n",
      ")\n",
      "STEP:  0\n",
      "loss: 0.5023738122475573\n",
      "loss: 0.49856639379435636\n",
      "loss: 0.47901196061152895\n",
      "loss: 0.4463349021484231\n",
      "loss: 0.3540631025749306\n",
      "loss: 0.20507016617681456\n",
      "loss: 1.3960531561165743\n",
      "loss: 0.032494411484716865\n",
      "loss: 0.029934875839604042\n",
      "loss: 0.028326821011534116\n",
      "loss: 0.026830612218822856\n",
      "loss: 0.02377120198998914\n",
      "loss: 0.0189014135045447\n",
      "loss: 0.010646818233204867\n",
      "loss: 0.008725752090268454\n",
      "loss: 0.007872181287777282\n",
      "loss: 0.005477842749594329\n",
      "loss: 0.004051933564063783\n",
      "loss: 0.0027296227011584305\n",
      "loss: 0.0015402652769799664\n",
      "test loss: 0.0013000876156952295\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import matplotlib\n",
    "# Non-interactive backend, you can't call plt.show() to see the figure interactively\n",
    "# matplotlib.use('Agg') must be placed before import matplotlib.pyplot\n",
    "matplotlib.use('Agg')\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def generateSineWave():\n",
    "    np.random.seed(2)\n",
    "    T = 20\n",
    "    L = 1000\n",
    "    N = 100\n",
    "    x = np.empty((N, L), 'int64') #the dataset has 100 items and each item's length is 1000\n",
    "    x[:] = np.array(range(L)) + np.random.randint(-4 * T, 4 * T, N).reshape(N, 1)\n",
    "    data = np.sin(x / 1.0 / T).astype('float64')\n",
    "    torch.save(data, open('traindata.pt', 'wb'))\n",
    "\n",
    "class Sequence(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Sequence, self).__init__()\n",
    "        self.lstm1 = nn.LSTMCell(1, 51)\n",
    "        self.lstm2 = nn.LSTMCell(51, 51)\n",
    "        self.linear = nn.Linear(51, 1)\n",
    "\n",
    "    def forward(self, input, future = 0):\n",
    "        outputs = []\n",
    "        h_t = torch.zeros(input.size(0), 51, dtype=torch.double)\n",
    "        c_t = torch.zeros(input.size(0), 51, dtype=torch.double)\n",
    "        h_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)\n",
    "        c_t2 = torch.zeros(input.size(0), 51, dtype=torch.double)\n",
    "\n",
    "        h_t = h_t.to(device)\n",
    "        c_t = c_t.to(device)\n",
    "        h_t2 = h_t2.to(device)\n",
    "        c_t2 = c_t2.to(device)\n",
    "\n",
    "        for i, input_t in enumerate(input.chunk(input.size(1), dim=1)):\n",
    "            h_t, c_t = self.lstm1(input_t, (h_t, c_t))\n",
    "            h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))\n",
    "            output = self.linear(h_t2)  # output.shape:[batch,1]\n",
    "            outputs += [output] # outputs.shape:[[batch,1],...[batch,1]], list composed of n [batch,1],\n",
    "        for i in range(future):# if we should predict the future\n",
    "            h_t, c_t = self.lstm1(output, (h_t, c_t))\n",
    "            h_t2, c_t2 = self.lstm2(h_t, (h_t2, c_t2))\n",
    "            output = self.linear(h_t2) # output.shape:[batch,1]\n",
    "            outputs += [output]  # outputs.shape:[[batch,1],...[batch,1]], list composed of n [batch,1],\n",
    "        outputs = torch.stack(outputs, 1).squeeze(2) # shape after stack:[batch, n, 1], shape after squeeze: [batch,n]\n",
    "        return outputs\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    # 1. generate sine wave data\n",
    "    generateSineWave()\n",
    "    # set random seed to 0\n",
    "    np.random.seed(0)\n",
    "    torch.manual_seed(0)\n",
    "    # load data and make training set\n",
    "    data = torch.load('traindata.pt')\n",
    "    input = torch.from_numpy(data[3:, :-1])\n",
    "    target = torch.from_numpy(data[3:, 1:])\n",
    "    test_input = torch.from_numpy(data[:3, :-1])\n",
    "    test_target = torch.from_numpy(data[:3, 1:])\n",
    "    input = input.to(device)\n",
    "    target = target.to(device)\n",
    "    test_input = test_input.to(device)\n",
    "    test_target = test_target.to(device)\n",
    "    # 2. build the model\n",
    "    seq = Sequence()\n",
    "    seq.double()\n",
    "    print(seq)\n",
    "    # move to cuda\n",
    "    # if torch.cuda.device_count()>1:\n",
    "    #     seq = nn.DataParallel(seq)\n",
    "    seq = seq.to(device)\n",
    "\n",
    "    # 3 loss function\n",
    "    criterion = nn.MSELoss()\n",
    "    # 4 use LBFGS as optimizer since we can load the whole data to train\n",
    "    optimizer = optim.LBFGS(seq.parameters(), lr=0.8)\n",
    "    # 5 begin to train\n",
    "    for i in range(1):\n",
    "        print('STEP: ', i)\n",
    "        def closure():\n",
    "            # forward\n",
    "            out = seq(input)\n",
    "            loss = criterion(out, target)\n",
    "            print('loss:', loss.item())\n",
    "            # backward\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            return loss\n",
    "        optimizer.step(closure)\n",
    "        # begin to predict, no need to track gradient here\n",
    "        with torch.no_grad():\n",
    "            future = 1000\n",
    "            pred = seq(test_input, future=future)\n",
    "            loss = criterion(pred[:, :-future], test_target)\n",
    "            print('test loss:', loss.item())\n",
    "\n",
    "            y = pred.detach().cpu()\n",
    "            y = y.numpy()\n",
    "        # draw the result\n",
    "        plt.figure(figsize=(30,10))\n",
    "        plt.title('Predict future values for time sequences\\n(Dashlines are predicted values)', fontsize=30)\n",
    "        plt.xlabel('x', fontsize=20)\n",
    "        plt.ylabel('y', fontsize=20)\n",
    "        plt.xticks(fontsize=20)\n",
    "        plt.yticks(fontsize=20)\n",
    "        def draw(yi, color):\n",
    "            plt.plot(np.arange(input.size(1)), yi[:input.size(1)], color, linewidth = 2.0)\n",
    "            plt.plot(np.arange(input.size(1), input.size(1) + future), yi[input.size(1):], color + ':', linewidth = 2.0)\n",
    "        draw(y[0], 'r')\n",
    "        draw(y[1], 'g')\n",
    "        draw(y[2], 'b')\n",
    "        plt.savefig('predict%d.pdf'%i)\n",
    "        plt.close()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}